from helper import *
from config import *
import numpy as np
import scipy
import scipy.signal
from bokeh.plotting import figure, output_file, show
from bokeh.models import (
    ColumnDataSource, Range1d, DataRange1d, DatetimeAxis,
    TickFormatter, DatetimeTickFormatter, FuncTickFormatter,
    Grid, Legend, Plot, BoxAnnotation, Span, CustomJS, Rect, Circle, Line,
    HoverTool, BoxZoomTool, PanTool, WheelZoomTool, ResetTool, SaveTool,
    WMTSTileSource, GMapPlot, GMapOptions,
    LabelSet, Label, ColorBar, LinearColorMapper, BasicTicker, PrintfTickFormatter
)

TOOLS = "pan,wheel_zoom,box_zoom,reset,save"
ACTIVE_SCROLL_TOOLS = "wheel_zoom"

ulog = load_ulog_file('1.ulg')
px4_ulog = PX4ULog(ulog)
data = ulog.data_list
cur_dataset = [elem for elem in data
               if elem.name == 'sensor_combined' and elem.multi_id == 0][0]

def expand_field_names(field_names, data_set):
    """
    expand field names if they're a function
    """
    field_names_expanded = []
    for field_name in field_names:
        if hasattr(field_name, '__call__'):
            new_field_name, new_data = field_name(cur_dataset.data)
            data_set[new_field_name] = new_data
            field_names_expanded.append(new_field_name)
        else:
            data_set[field_name] = cur_dataset.data[field_name]
            field_names_expanded.append(field_name)
    return field_names_expanded


data_set = {}
timestamp_key = 'timestamp'  # 微秒为单位
if 'timestamp_sample' in cur_dataset.data.keys():
    timestamp_key = 'timestamp_sample'
data_set[timestamp_key] = cur_dataset.data[timestamp_key]

# 计算采样频率，这里不考虑日志的丢失
delta_t = ((data_set[timestamp_key][-1] - data_set[timestamp_key][0]) * 1.0e-6) / len(data_set[timestamp_key])
# print(delta_t, data_set[timestamp_key][-1], data_set[timestamp_key][0], len(data_set[timestamp_key]))
if delta_t < 0.000001:  # 避免除以0
    print("error")
sampling_frequency = int(1.0 / delta_t)
# print(sampling_frequency)
if sampling_frequency < 100:  # 要求满足最小的采样频率100Hz
    print("error")

x_range_offset = (ulog.last_timestamp - ulog.start_timestamp) * 0.05
x_range = Range1d(ulog.start_timestamp - x_range_offset, ulog.last_timestamp + x_range_offset)

field_names_expanded = expand_field_names(['accelerometer_m_s2[0]', 'accelerometer_m_s2[1]', 'accelerometer_m_s2[2]'],
                                          data_set)
# print(field_names_expanded)

# calculate the spectrogram
psd = {}
for key in field_names_expanded:
    frequency, time, psd[key] = scipy.signal.spectrogram(
        data_set[key], fs=sampling_frequency, window='hann',
        nperseg=256, noverlap=128, scaling='density')

# sum all psd's
key_it = iter(psd)
sum_psd = psd[next(key_it)]
for key in key_it:
    sum_psd += psd[key]

# scale time to microseconds and add start time as offset
time = time * 1.0e6 + cur_dataset.data[timestamp_key][0]

image = [10 * np.log10(sum_psd)]
title = 'Acceleration Power Spectral Density'
for legend in ['X', 'Y', 'Z']:
    title += " " + legend
title += " [dB]"

# assume maximal data points per pixel at full resolution
max_num_data_points = 2.0 * plot_width
if len(time) > max_num_data_points:
    step_size = int(len(time) / max_num_data_points)
    time = time[::step_size]
    image[0] = image[0][:, ::step_size]

color_mapper = LinearColorMapper(palette="Viridis256", low=np.amin(image), high=np.amax(image))

p = figure(title='Acceleration Power Spectral Density', x_axis_label=None,
           y_axis_label='[Hz]', tools=TOOLS,
           active_scroll=ACTIVE_SCROLL_TOOLS)


p.xgrid.grid_line_color = None
p.ygrid.grid_line_color = 'navy'
p.ygrid.grid_line_alpha = 0.13
p.ygrid.minor_grid_line_color = 'navy'
p.ygrid.minor_grid_line_alpha = 0.05
p.plot_width = 840
p.plot_height = 336

p.x_range = Range1d(x_range.start, x_range.end)
p.y_range = Range1d(frequency[0], frequency[-1])
p.toolbar_location = 'above'
p.image(image=image, x=time[0], y=frequency[0], dw=(time[-1] - time[0]),
        dh=(frequency[-1] - frequency[0]), color_mapper=color_mapper)

color_bar = ColorBar(color_mapper=color_mapper,
                     major_label_text_font_size="5pt",
                     ticker=BasicTicker(desired_num_ticks=5),
                     formatter=PrintfTickFormatter(format="%f"),
                     title='[dB]',
                     label_standoff=6, border_line_color=None, location=(0, 0))
p.add_layout(color_bar, 'right')

# add plot zoom tool that only zooms in time axis
wheel_zoom = WheelZoomTool()
p.toolbar.tools = [PanTool(), wheel_zoom, BoxZoomTool(), ResetTool(), SaveTool()]  # updated_tools
p.toolbar.active_scroll = wheel_zoom

show(p)
